{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "State = (cars_loc1, cars_loc2) 每方最多20辆车\n",
    "\n",
    "行为 每晚最多移动-5～5辆车从loc1到loc2\n",
    "\n",
    "奖励 R 每租出去1辆车，奖励$10\n",
    "\n",
    "状态转移概率\n",
    "租车还车服从泊松分布\n",
    " - 第一个地点 平均租车3，返还3\n",
    " - 第二个地点 平均租车4， 返还2\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# 动态规划算法时间复杂度随着状态数量和行为数量的增多而成平方级别增多，\n",
    "# 本例中的状态数为21*21=441,行为数为11,按照动态规划算法完成一次迭代话费计算资源较昂贵\n",
    "# \n",
    "MAX_CARS = 20\n",
    "MAX_MOVE = 5\n",
    "S = []\n",
    "for i in range(MAX_CARS+1):\n",
    "    for j in range(MAX_CARS+1):\n",
    "        S.append((i,j))\n",
    "\n",
    "A = [i for i in range(-MAX_MOVE,MAX_MOVE+1)]\n",
    "params = {\n",
    "    \"loc1_rent\":3,\n",
    "    \"loc1_return\":3,\n",
    "    \"loc2_rent\":4,\n",
    "    \"loc2_return\":2\n",
    "}\n",
    "# print(S)\n",
    "# print(A)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def dynamics(s, a):\n",
    "    '''Jack Car 租赁问题的动力学，包含泊松分布等随机性\n",
    "    Args:\n",
    "        s 两地(loc1, loc2)汽车库存数 tuple (n1, n2) 0<=n1,n2<=20\n",
    "        a 转场汽车数 int 定义为从 loc1转运a辆车至loc2地，负数表示反向搬运 -5<=a<=5\n",
    "    Returns:\n",
    "        s_prime 后续状态 tuple(int, int) (n1,n2)\n",
    "        reward float 奖励，两地合算，每租出1辆车，奖励10\n",
    "        is_end Bool 本例可循环，没有终止状态，返回False\n",
    "    '''\n",
    "    s_prime, reward, is_end = None, None, False\n",
    "    # Add your code here\n",
    "    return s_prime, reward, is_end\n",
    "\n",
    "def poisson_prob(n, lamda):\n",
    "    '''计算从一个参数为n定义的泊松分布中，采样输出为lamda的概率\n",
    "    '''\n",
    "    if lamda < 0:\n",
    "        return 0.0\n",
    "    return np.power(lamda,n)*np.power(np.e, -lamda) / np.math.factorial(n)\n",
    "\n",
    "def _need_return(car_rent, loc_s0, loc_s1):\n",
    "    '''借出car_rent辆车，需要有多少辆车返回才能满足从loc_s0->loc_s1\n",
    "    '''\n",
    "    return car_rent + (loc_s1 - loc_s0)\n",
    "\n",
    "def prob_one_location(n_rent, n_return, loc_s0, loc_s1):\n",
    "    '''计算由借出和返还泊松参数决定的随机事件中，某一地点的汽车数量\n",
    "    从loc_s0变为loc_s1的概率。\n",
    "    '''\n",
    "    max_for_rent = loc_s0 # 当前最大可借出数量\n",
    "    prob = 0 # 初始概率设为0\n",
    "    for car_rent in range(0, max_for_rent):# 对于每一个借出的可能性\n",
    "        prob_rent = poisson_prob(n_rent, car_rent)\n",
    "        # 借出car_rent辆车，需要有car_return辆返回才能满足最终汽车数为loc_s1\n",
    "        car_need_return = _need_return(car_rent, loc_s0, loc_s1)\n",
    "        if car_need_return < 0: \n",
    "        # <0 说明借的不够多，将在借出里考虑计算，不重复计算\n",
    "            prob_return = 0.0\n",
    "            # continue\n",
    "        prob_return = poisson_prob(n_return, car_need_return)\n",
    "        prob += prob_rent * prob_return # 概率乘积\n",
    "    return prob\n",
    "    \n",
    "def reward_one_location(n_rent, loc_aftmv):\n",
    "    reward = 0\n",
    "    for cars in range(loc_aftmv + 1):\n",
    "        reward += cars * poisson_prob(n_rent, cars)\n",
    "    return reward"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "_need_return(3,5,2)\n",
    "prob_one_location(3,3,4,3)\n",
    "reward_one_location(3,5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def P(s, a, s1):\n",
    "    '''重新改写状态转换概率\n",
    "    '''\n",
    "    n1, n2 = s # 当前1,2两地汽车数量数目\n",
    "    if not(0<= n1 <= MAX_CARS and 0<= n2 <= MAX_CARS):\n",
    "        #\"初始状态不合法\"\n",
    "        return 0.0\n",
    "    n1_aftmv, n2_aftmv = n1 - a, n2 + a # 移动a辆汽车后两地汽车数量\n",
    "    if not(0<= n1_aftmv <= MAX_CARS and 0 <= n2_aftmv <= MAX_CARS): \n",
    "        # \"转场的汽车数不符合要求\"\n",
    "        return 0.0\n",
    "    n1_prime, n2_prime = s1\n",
    "    prob = prob_one_location(params[\"loc1_rent\"],\n",
    "                             params[\"loc1_return\"],\n",
    "                             n1_aftmv, n1_prime)\n",
    "    prob *= prob_one_location(params[\"loc2_rent\"],\n",
    "                              params[\"loc2_return\"],\n",
    "                              n2_aftmv, n2_prime)\n",
    "    return prob\n",
    "\n",
    "\n",
    "def R(s, a):\n",
    "    n1, n2 = s # 当前1,2两地汽车数量数目\n",
    "    if not(0<= n1 <= MAX_CARS and 0<= n2 <= MAX_CARS):\n",
    "        #\"初始状态不合法\"\n",
    "        return 0.0\n",
    "    n1_aftmv, n2_aftmv = n1 - a, n2 + a # 移动a辆汽车后两地汽车数量\n",
    "    if not(0<= n1_aftmv <= MAX_CARS and 0 <= n2_aftmv <= MAX_CARS): \n",
    "        # \"转场的汽车数不符合要求\"\n",
    "        return 0.0\n",
    "    reward = 0\n",
    "    reward += reward_one_location(params[\"loc1_rent\"], n1_aftmv)\n",
    "    reward += reward_one_location(params[\"loc2_rent\"], n2_aftmv)\n",
    "    return reward\n",
    "    \n",
    "def set_value(V, s, v):\n",
    "    loc1, loc2 = s\n",
    "    V[loc1,loc2] = v\n",
    "    \n",
    "def get_value(V, s):\n",
    "    loc1, loc2 = s\n",
    "    return V[loc1, loc2]\n",
    "\n",
    "def display_V(V):\n",
    "    print(V)\n",
    "    \n",
    "gamma = 1\n",
    "MDP = S, A, R, P, gamma   \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#from utils import get_prob, get_reward\n",
    "#from utils import get_pi, uniform_random_pi, greedy_pi\n",
    "#from utils import policy_evaluate, policy_iterate, value_iterate\n",
    "from tqdm import tqdm\n",
    "'''\n",
    "def P(s, a, s1):\n",
    "    s_prime, _, _ = dynamics(s, a)\n",
    "    return s1 == s_prime\n",
    "\n",
    "def R(s, a):\n",
    "    _, r, _ = dynamics(s, a)\n",
    "    return r\n",
    "'''\n",
    "def get_prob(P, s, a, s1):\n",
    "    return P(s, a, s1)\n",
    "\n",
    "def get_reward(R, s, a):\n",
    "    return R(s, a)\n",
    "\n",
    "########################\n",
    "\n",
    "def get_pi(Pi, s, a, MDP = None, V = None):\n",
    "    return Pi(MDP, V, s, a)\n",
    "\n",
    "def uniform_random_pi(MDP = None, V = None, s = None, a = None):\n",
    "    _, A, _, _, _ = MDP\n",
    "    n = len(A)\n",
    "    return 0 if n == 0 else 1.0/n\n",
    "\n",
    "def greedy_pi(MDP, V, s, a):\n",
    "    S, A, P, R, gamma = MDP\n",
    "    max_v, a_max_v = -float('inf'), []\n",
    "    for a_opt in A:# 统计后续状态的最大价值以及到达到达该状态的行为（可能不止一个）\n",
    "        s_prime, reward, _ = dynamics(s, a_opt)\n",
    "        v_s_prime = get_value(V, s_prime)\n",
    "        if v_s_prime > max_v:\n",
    "            max_v = v_s_prime\n",
    "            a_max_v = [a_opt]\n",
    "        elif(v_s_prime == max_v):\n",
    "            a_max_v.append(a_opt)\n",
    "    n = len(a_max_v)\n",
    "    if n == 0: return 0.0\n",
    "    return 1.0/n if a in a_max_v else 0.0\n",
    "\n",
    "\n",
    "def epsilon_greedy_pi(MDP, V, s, a, epsilon = 0.1):\n",
    "    if MDP is None:\n",
    "        return 0.0\n",
    "    _, A, _, _, _ = MDP\n",
    "    m = len(A)\n",
    "    greedy_p = greedy_pi(MDP, V, s, a)\n",
    "    if greedy_p == 0:\n",
    "        return epsilon / m\n",
    "    # n = int(1.0/greedy_p)\n",
    "    return (1 - epsilon + epsilon/m) * greedy_p\n",
    "\n",
    "\n",
    "def compute_q(MDP, V, s, a):\n",
    "    '''根据给定的MDP，价值函数V，计算状态行为对s,a的价值qsa\n",
    "    '''\n",
    "    S, A, R, P, gamma = MDP\n",
    "    q_sa = 0\n",
    "    for s_prime in S:\n",
    "        q_sa += get_prob(P, s, a, s_prime) * get_value(V, s_prime)\n",
    "    q_sa = get_reward(R, s,a) + gamma * q_sa\n",
    "    return q_sa\n",
    "\n",
    "\n",
    "def compute_v(MDP, V, Pi, s):\n",
    "    '''给定MDP下依据某一策略Pi和当前状态价值函数V计算某状态s的价值\n",
    "    '''\n",
    "    S, A, R, P, gamma = MDP\n",
    "    v_s = 0\n",
    "    for a in A:\n",
    "        v_s += get_pi(Pi, s, a, MDP, V) * compute_q(MDP, V, s, a)\n",
    "    return v_s        \n",
    "\n",
    "def update_V(MDP, V, Pi):\n",
    "    '''给定一个MDP和一个策略，更新该策略下的价值函数V\n",
    "    '''\n",
    "    S, _, _, _, _ = MDP\n",
    "    V_prime = V.copy()\n",
    "    for s in S:\n",
    "        #print(\"in update V, for state s\")\n",
    "        set_value(V_prime, s, compute_v(MDP, V_prime, Pi, s))\n",
    "    return V_prime\n",
    "\n",
    "\n",
    "def policy_evaluate(MDP, V, Pi, n):\n",
    "    '''使用n次迭代计算来评估一个MDP在给定策略Pi下的状态价值，初始时价值为V\n",
    "    '''\n",
    "    for i in tqdm(range(n)):\n",
    "        #print(\"====第{}次迭代====\".format(i+1))\n",
    "        V = update_V(MDP, V, Pi)\n",
    "        #display_V(V)\n",
    "    return V\n",
    "\n",
    "def policy_iterate(MDP, V, Pi, n, m):\n",
    "    cur_Pi = Pi\n",
    "    for i in range(m):\n",
    "        V = policy_evaluate(MDP, V, Pi, n)\n",
    "        Pi = epsilon_greedy_pi\n",
    "        #print(\"改善了策略\")\n",
    "    return V\n",
    "\n",
    "# 价值迭代得到最优状态价值过程\n",
    "def compute_v_from_max_q(MDP, V, s):\n",
    "    '''根据一个状态的下所有可能的行为价值中最大一个来确定当前状态价值\n",
    "    '''\n",
    "    S, A, R, P, gamma = MDP\n",
    "    v_s = -float('inf')\n",
    "    for a in A:\n",
    "        qsa = compute_q(MDP, V, s, a)\n",
    "        if qsa >= v_s:\n",
    "            v_s = qsa\n",
    "    return v_s\n",
    "\n",
    "def update_V_without_pi(MDP, V):\n",
    "    '''在不依赖策略的情况下直接通过后续状态的价值来更新状态价值\n",
    "    '''\n",
    "    S, _, _, _, _ = MDP\n",
    "    V_prime = V.copy()\n",
    "    for s in S:\n",
    "        set_value(V_prime, s, compute_v_from_max_q(MDP, V_prime, s))\n",
    "    return V_prime\n",
    "\n",
    "def value_iterate(MDP, V, n):\n",
    "    '''价值迭代\n",
    "    '''\n",
    "    for i in range(n):\n",
    "        V = update_V_without_pi(MDP, V)\n",
    "        display_V(V)\n",
    "    return V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "V = np.zeros((MAX_CARS+1, MAX_CARS+1))\n",
    "V_pi = value_iterate(MDP, V, 4)\n",
    "#V_pi = policy_evaluate(MDP, V, uniform_random_pi, 1)\n",
    "display_V(V_pi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "plt.imshow(V_pi, cmap=plt.cm.cool, interpolation=None, origin=\"lower\")#, extent=[0, 11, 0, 22])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "v_a = np.range((MAX_CARS+1,MAX_CARS+1))\n",
    "for i in range(MAX_CARS+1):\n",
    "    for j in range(MAX_CARS+1):\n",
    "        v_a[i,j] = greedy_policy"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.5.2"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
